package com.pai4j.zwiki.service.agent;

import lombok.extern.slf4j.Slf4j;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.concurrent.TimeUnit;
import java.util.stream.Collectors;

/**
 * 代码上下文管理器
 * 
 * 核心功能：
 * 1. 管理LLM分析过程中的上下文状态
 * 2. 实现上下文压缩和摘要
 * 3. 上下文记忆机制（短期记忆 + 长期记忆）
 * 4. 智能上下文检索和关联
 * 5. 上下文持久化到Redis
 *
 * 设计理念：
 * - 短期记忆：最近3-5轮分析的完整上下文
 * - 长期记忆：关键信息的压缩摘要
 * - 检索机制：根据当前任务检索相关的历史上下文
 *
 * @author pai4j
 * @since 2025-10-14
 */
@Slf4j
@Service
public class CodeContextManager {

    private final RedisTemplate<String, Object> redisTemplate;

    // 上下文窗口大小限制
    private static final int MAX_CONTEXT_SIZE = 50000; // 字符数
    private static final int SHORT_TERM_MEMORY_SIZE = 5; // 保留最近5轮
    private static final int LONG_TERM_MEMORY_SIZE = 20; // 保留20条关键摘要

    public CodeContextManager(RedisTemplate<String, Object> redisTemplate) {
        this.redisTemplate = redisTemplate;
    }

    /**
     * 创建新的上下文会话
     */
    public String createContextSession(String taskId) {
        String sessionId = taskId + "_context_" + System.currentTimeMillis();
        
        Map<String, Object> session = new HashMap<>();
        session.put("sessionId", sessionId);
        session.put("taskId", taskId);
        session.put("createdAt", System.currentTimeMillis());
        session.put("shortTermMemory", new ArrayList<>());
        session.put("longTermMemory", new ArrayList<>());
        session.put("roundCount", 0);

        // 存储到Redis（24小时过期）
        String key = "zwiki:context:session:" + sessionId;
        redisTemplate.opsForValue().set(key, session, 24, TimeUnit.HOURS);

        log.info("创建上下文会话: sessionId={}", sessionId);
        return sessionId;
    }

    /**
     * 添加分析结果到上下文
     */
    public void addAnalysisToContext(String sessionId, String roundName, 
                                     String analysis, int importance) {
        try {
            String key = "zwiki:context:session:" + sessionId;
            @SuppressWarnings("unchecked")
            Map<String, Object> session = (Map<String, Object>) redisTemplate.opsForValue().get(key);

            if (session == null) {
                log.warn("上下文会话不存在: {}", sessionId);
                return;
            }

            // 添加到短期记忆
            @SuppressWarnings("unchecked")
            List<Map<String, Object>> shortTermMemory = 
                    (List<Map<String, Object>>) session.get("shortTermMemory");
            
            Map<String, Object> memory = new HashMap<>();
            memory.put("round", roundName);
            memory.put("analysis", analysis);
            memory.put("importance", importance);
            memory.put("timestamp", System.currentTimeMillis());
            
            shortTermMemory.add(memory);

            // 保持短期记忆大小
            if (shortTermMemory.size() > SHORT_TERM_MEMORY_SIZE) {
                shortTermMemory.remove(0);
            }

            // 如果重要度高，添加到长期记忆（压缩版本）
            if (importance >= 8) {
                @SuppressWarnings("unchecked")
                List<Map<String, Object>> longTermMemory = 
                        (List<Map<String, Object>>) session.get("longTermMemory");
                
                Map<String, Object> longMemory = new HashMap<>();
                longMemory.put("round", roundName);
                longMemory.put("summary", compressContext(analysis, 500));
                longMemory.put("importance", importance);
                
                longTermMemory.add(longMemory);

                // 保持长期记忆大小
                if (longTermMemory.size() > LONG_TERM_MEMORY_SIZE) {
                    // 移除重要度最低的
                    longTermMemory.sort(Comparator.comparingInt(
                            m -> (Integer) m.get("importance")));
                    longTermMemory.remove(0);
                }
            }

            // 更新轮次计数
            session.put("roundCount", ((Integer) session.get("roundCount")) + 1);

            // 保存回Redis
            redisTemplate.opsForValue().set(key, session, 24, TimeUnit.HOURS);

            log.debug("添加分析到上下文: sessionId={}, round={}, importance={}", 
                    sessionId, roundName, importance);

        } catch (Exception e) {
            log.error("添加分析到上下文失败", e);
        }
    }

    /**
     * 获取当前完整上下文（用于传递给LLM）
     */
    public String getFullContext(String sessionId) {
        try {
            String key = "zwiki:context:session:" + sessionId;
            @SuppressWarnings("unchecked")
            Map<String, Object> session = (Map<String, Object>) redisTemplate.opsForValue().get(key);

            if (session == null) {
                return "";
            }

            StringBuilder context = new StringBuilder();

            // 添加长期记忆摘要
            @SuppressWarnings("unchecked")
            List<Map<String, Object>> longTermMemory = 
                    (List<Map<String, Object>>) session.get("longTermMemory");
            
            if (longTermMemory != null && !longTermMemory.isEmpty()) {
                context.append("# 关键理解（长期记忆）\n\n");
                longTermMemory.forEach(mem -> {
                    context.append("## ").append(mem.get("round")).append("\n");
                    context.append(mem.get("summary")).append("\n\n");
                });
            }

            // 添加短期记忆（最近几轮的完整分析）
            @SuppressWarnings("unchecked")
            List<Map<String, Object>> shortTermMemory = 
                    (List<Map<String, Object>>) session.get("shortTermMemory");
            
            if (shortTermMemory != null && !shortTermMemory.isEmpty()) {
                context.append("# 最近分析（短期记忆）\n\n");
                shortTermMemory.forEach(mem -> {
                    context.append("## ").append(mem.get("round")).append("\n");
                    String analysis = (String) mem.get("analysis");
                    // 短期记忆可以更详细，但也要控制大小
                    context.append(compressContext(analysis, 2000)).append("\n\n");
                });
            }

            String fullContext = context.toString();
            
            // 如果上下文过大，进行智能压缩
            if (fullContext.length() > MAX_CONTEXT_SIZE) {
                fullContext = intelligentCompress(fullContext, MAX_CONTEXT_SIZE);
                log.info("上下文过大，已压缩: {} → {} 字符", 
                        context.length(), fullContext.length());
            }

            return fullContext;

        } catch (Exception e) {
            log.error("获取上下文失败", e);
            return "";
        }
    }

    /**
     * 检索相关上下文（基于关键词）
     */
    public String retrieveRelevantContext(String sessionId, String query) {
        try {
            String key = "zwiki:context:session:" + sessionId;
            @SuppressWarnings("unchecked")
            Map<String, Object> session = (Map<String, Object>) redisTemplate.opsForValue().get(key);

            if (session == null) {
                return "";
            }

            // 提取查询关键词
            Set<String> queryKeywords = extractKeywords(query);

            // 从短期记忆中检索
            @SuppressWarnings("unchecked")
            List<Map<String, Object>> shortTermMemory = 
                    (List<Map<String, Object>>) session.get("shortTermMemory");

            List<Map<String, Object>> relevantMemories = shortTermMemory.stream()
                    .filter(mem -> isRelevant((String) mem.get("analysis"), queryKeywords))
                    .collect(Collectors.toList());

            // 构建相关上下文
            StringBuilder context = new StringBuilder();
            context.append("# 相关上下文\n\n");

            relevantMemories.forEach(mem -> {
                context.append("## ").append(mem.get("round")).append("\n");
                context.append(mem.get("analysis")).append("\n\n");
            });

            return compressContext(context.toString(), 10000);

        } catch (Exception e) {
            log.error("检索相关上下文失败", e);
            return "";
        }
    }

    /**
     * 智能压缩上下文
     */
    private String intelligentCompress(String context, int maxSize) {
        if (context.length() <= maxSize) {
            return context;
        }

        // 策略：保留标题和关键信息，压缩详细内容
        String[] sections = context.split("##");
        StringBuilder compressed = new StringBuilder();
        int currentSize = 0;

        for (String section : sections) {
            if (section.trim().isEmpty()) continue;

            // 保留标题
            String[] lines = section.split("\n", 2);
            String title = "## " + (lines.length > 0 ? lines[0] : "");
            String content = lines.length > 1 ? lines[1] : "";

            compressed.append(title).append("\n");
            currentSize += title.length();

            // 压缩内容
            if (currentSize + content.length() <= maxSize) {
                compressed.append(content).append("\n");
                currentSize += content.length();
            } else {
                int remainingSpace = maxSize - currentSize - 100;
                if (remainingSpace > 100) {
                    compressed.append(compressContext(content, remainingSpace)).append("\n");
                }
                break;
            }
        }

        return compressed.toString();
    }

    /**
     * 压缩上下文到指定长度
     */
    private String compressContext(String text, int maxLength) {
        if (text == null) return "";
        if (text.length() <= maxLength) return text;

        // 提取关键句子
        String[] sentences = text.split("[。！？\\n]");
        List<String> importantSentences = Arrays.stream(sentences)
                .filter(s -> !s.trim().isEmpty())
                .filter(this::isImportantSentence)
                .limit(maxLength / 50)
                .collect(Collectors.toList());

        String compressed = String.join("。", importantSentences);
        
        if (compressed.length() > maxLength) {
            compressed = compressed.substring(0, maxLength) + "...";
        }

        return compressed;
    }

    /**
     * 判断是否为重要句子
     */
    private boolean isImportantSentence(String sentence) {
        String lower = sentence.toLowerCase();
        // 包含关键词的句子更重要
        return lower.contains("核心") || lower.contains("关键") || 
               lower.contains("重要") || lower.contains("主要") ||
               lower.contains("架构") || lower.contains("设计") ||
               lower.contains("问题") || lower.contains("建议");
    }

    /**
     * 提取关键词
     */
    private Set<String> extractKeywords(String text) {
        Set<String> keywords = new HashSet<>();
        String[] words = text.split("\\s+");
        
        for (String word : words) {
            if (word.length() >= 2 && !isStopWord(word)) {
                keywords.add(word.toLowerCase());
            }
        }

        return keywords;
    }

    /**
     * 判断是否为停用词
     */
    private boolean isStopWord(String word) {
        Set<String> stopWords = Set.of("的", "是", "在", "了", "和", "有", "为", "与");
        return stopWords.contains(word);
    }

    /**
     * 判断文本是否与关键词相关
     */
    private boolean isRelevant(String text, Set<String> keywords) {
        String lowerText = text.toLowerCase();
        return keywords.stream().anyMatch(lowerText::contains);
    }

    /**
     * 清除会话上下文
     */
    public void clearSession(String sessionId) {
        String key = "zwiki:context:session:" + sessionId;
        redisTemplate.delete(key);
        log.info("清除上下文会话: sessionId={}", sessionId);
    }
}

